"""
增强多信号策略 - 整合 RSRS、OC_DIFF、RET 三个因子
"""
import numpy as np
import pandas as pd
from sklearn.linear_model import LinearRegression
from .base_strategy import BaseStrategy


class MultiSignalStrategy(BaseStrategy):
    strategy_name = "multi_signal"
    
    def __init__(self, api, symbol_info, market_period):
        super().__init__(api, symbol_info, market_period)
        for symbol in symbol_info.keys():
            self.all_info[symbol]["rsrs_slopes"] = []
            self.all_info[symbol]["corr_series"] = []
            self.all_info[symbol]["ret_series"] = []
    
    def calculate_signal(self, symbol, current_pos, kline):
        """计算综合信号 - 整合 RSRS、OC_DIFF、RET 三个因子"""
        # 计算三个因子的得分
        rsrs_score = self._calculate_rsrs(symbol, kline)
        oc_score = self._calculate_oc_corr(symbol, kline)
        ret_score = self._calculate_ret(symbol, kline)
        
        # 收集所有因子得分
        all_scores = [rsrs_score, oc_score, ret_score]
        
        # 综合得分（等权重平均）
        combined_score = np.mean(all_scores)
        
        # 获取阈值参数
        long_threshold = self.all_info[symbol].get("long_threshold", 0.5)
        short_threshold = self.all_info[symbol].get("short_threshold", -0.5)
        close_threshold = self.all_info[symbol].get("close_threshold", 0.1)
        
        # 根据综合得分决定信号
        if current_pos == 0:
            if combined_score > long_threshold:
                return 1, combined_score
            elif combined_score < short_threshold:
                return -1, combined_score
            else:
                return 0, combined_score
        elif current_pos > 0:
            if combined_score < short_threshold:
                return -1, combined_score
            elif combined_score < -close_threshold:
                return 0, combined_score
            else:
                return 1, combined_score
        elif current_pos < 0:
            if combined_score > long_threshold:
                return 1, combined_score
            elif combined_score > close_threshold:
                return 0, combined_score
            else:
                return -1, combined_score
    
    # ==================== 因子计算方法 ====================
    
    def _calculate_rsrs(self, symbol, kline):
        """计算 RSRS 因子得分"""
        rsrs_period = self.all_info[symbol].get("rsrs_period", 18)
        rsrs_std_period = self.all_info[symbol].get("rsrs_std_period", 600)
        
        high_prices = kline.high.values
        low_prices = kline.low.values
        rsrs_slopes = self.all_info[symbol]["rsrs_slopes"]
        
        current_len = len(high_prices)
        if len(rsrs_slopes) < current_len - rsrs_period:
            start_idx = max(rsrs_period, len(rsrs_slopes))
            for i in range(start_idx, min(current_len, rsrs_std_period + rsrs_period)):
                window_high = high_prices[i-rsrs_period:i]
                window_low = low_prices[i-rsrs_period:i]
                X = window_low.reshape(-1, 1)
                y = window_high
                reg = LinearRegression().fit(X, y)
                self.all_info[symbol]["rsrs_slopes"].append(reg.coef_[0])
        
        if len(rsrs_slopes) >= rsrs_std_period:
            window_high = high_prices[-rsrs_period:]
            window_low = low_prices[-rsrs_period:]
            X = window_low.reshape(-1, 1)
            y = window_high
            reg = LinearRegression().fit(X, y)
            slope = reg.coef_[0]
            self.all_info[symbol]["rsrs_slopes"][:-1] = self.all_info[symbol]["rsrs_slopes"][1:]
            self.all_info[symbol]["rsrs_slopes"][-1] = slope
        
        if len(rsrs_slopes) == 0:
            return 0.0
            
        current_slope = self.all_info[symbol]["rsrs_slopes"][-1]
        recent_slopes = self.all_info[symbol]["rsrs_slopes"][-rsrs_std_period:]
        mean_slope = np.mean(recent_slopes)
        std_slope = np.std(recent_slopes)
        
        return (current_slope - mean_slope) / std_slope if std_slope > 0 else 0.0
    
    def _calculate_oc_corr(self, symbol, kline):
        """计算 OC 相关性因子得分"""
        score_period = self.all_info[symbol].get("oc_score_period", 20)
        std_period = self.all_info[symbol].get("oc_std_period", 600)
        
        close_prices = kline.close.values
        open_prices = kline.open.values
        corr_series = self.all_info[symbol]["corr_series"]
        
        current_len = len(close_prices)
        if len(corr_series) < current_len - score_period:
            start_idx = max(score_period, len(corr_series))
            for i in range(start_idx, min(current_len, std_period + score_period)):
                window_close = close_prices[i-score_period:i]
                window_open = open_prices[i-score_period:i]
                mem = (window_close - window_open).mean()
                corr = np.corrcoef(window_close, window_open)[0][1] * (1 if mem >= 0 else -1)
                self.all_info[symbol]["corr_series"].append(corr)
        
        if len(corr_series) >= std_period:
            window_close = close_prices[-score_period:]
            window_open = open_prices[-score_period:]
            mem = (window_close - window_open).mean()
            corr = np.corrcoef(window_close, window_open)[0][1] * (1 if mem >= 0 else -1)
            self.all_info[symbol]["corr_series"][:-1] = self.all_info[symbol]["corr_series"][1:]
            self.all_info[symbol]["corr_series"][-1] = corr
        
        if len(corr_series) == 0:
            return 0.0
            
        current_score = self.all_info[symbol]["corr_series"][-1]
        recent_score = self.all_info[symbol]["corr_series"][-std_period:]
        mean_score = np.mean(recent_score)
        std_score = np.std(recent_score)
        
        return (current_score - mean_score) / std_score if std_score > 0 else 0.0
    
    def _calculate_ret(self, symbol, kline):
        """计算 RET 因子得分"""
        open_prices = kline.open.values
        high_prices = kline.high.values
        low_prices = kline.low.values
        close_prices = kline.close.values
        
        ret_series = self.all_info[symbol]["ret_series"]
        
        if len(ret_series) == 0:
            self.all_info[symbol]["ret_series"] = list(high_prices + low_prices - open_prices - close_prices)
        else:
            ret_value = high_prices[-1] + low_prices[-1] - open_prices[-1] - close_prices[-1]
            self.all_info[symbol]["ret_series"].append(ret_value)
            if len(self.all_info[symbol]["ret_series"]) > 800:
                self.all_info[symbol]["ret_series"] = self.all_info[symbol]["ret_series"][-800:]
        
        return self.all_info[symbol]["ret_series"][-1] if len(self.all_info[symbol]["ret_series"]) > 0 else 0.0
    

